# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/28 11:07
# @Function:
"""
水源涵养：除了海洋生态系统都要计算；
最简单的计算方法：

输入：某区域的生态图.tif。
判断核算的区域是否有海洋：
    有则：
        将生态图分成有海洋和无海洋生态图；
    否则：
        用输入的生态图直接计算；

由于所有的生态类型都用一个公式，参数也统一，所以直接将生态图做成布尔数组，无效值为fasle，有值则为true；

A：（面积：0.0009）
P：（降雨量：tif矩阵）
R：（径流量：tif矩阵）
ET：（蒸散发量：tif矩阵）
result = 生态图布尔矩阵 * A * （P - R - ET） * 1000

将计算结果保存成tif
"""

import os
import re
# import time
from concurrent.futures import ThreadPoolExecutor, as_completed, ProcessPoolExecutor
from typing import List, Any
import numpy as np
from tqdm import tqdm
# import pandas as pd
from calc_v1 import Formula, MatrixCalculator
from calc_v1.formulas import formulas_list_Q
from calc_v1.utils import load_parameters, load_parameters_from_json, calculate_with_excel_mapping, \
    calculate_with_boolean_mapping, replace_tif_data_with_array


def run_matrix_calculator(
        json_file_path: str, mode: str, formulas: List[Formula] = None,
        output_base_path: str = r"F:\code\dev\calc-gep-regulate-cqc\calc_v1\output2") -> dict[Any, int | Any] | None:
    """
    运行矩阵计算器，加载参数并执行指定的公式。

    参数:
    json_file_path (str): 包含参数的 JSON 文件路径。
    mode (str): 当前计算模式。
    formulas (List[Formula]): 可用的公式列表。默认值为 formulas_list_Q。

    返回:
    None

    异常:
    Exception: 如果加载文件或执行公式时出错，会打印错误信息。

    程序逻辑：
    1. 初始化计算器
    2. 加载参数
    3. 判断可用公式 （如果提示没有可运行的公式，应该提示缺少哪些变量）
    4. 计算
    4.1 获取11个细分的生态图 （写一个函数，输入一个tif文件，返回一个数据列表）extract_ecosystem_arrays（tif_path）
        json_data = load_parameters_from_json(json_file_path)
        tif_path = json_data['eco_tif']['file_path']
        ecosystem_arrays = extract_ecosystem_arrays(tif_path)
    4.2

    """
    if formulas is None:  # 如果未传递 formulas，则加载默认值
        from calc_v1.formulas import formulas_list_Q
        formulas = formulas_list_Q

    # 初始化计算器
    global extracted_part
    calculator = MatrixCalculator()
    calculator.load_formulas(formulas)
    #  修改了mode的用法，此处无意义
    # calculator.set_mode(mode)

    # 加载参数
    parameters = load_parameters(json_file_path, mode)
    for parameter in parameters:
        calculator.param_manager.add_parameter(parameter)

    # 获取可用公式
    available_formulas, formula_excel_name = calculator.prompt_available_formulas()

    # 提示用户可以选择运行的公式
    if available_formulas:
        print("可以运行以下公式：")
        for formula in available_formulas:
            print(f"- {formula}")
    else:
        print("没有可运行的公式。")
        return

    # todo 需要加一层循环，9个生态图
    json_data = load_parameters_from_json(json_file_path)
    tif_path = json_data['eco_tif']['file_path']
    print("eco_tif_path：", tif_path)
    # 提取文件名
    file_name = os.path.basename(tif_path)

    # 使用正则表达式提取数字部分
    match = re.search(r'(\d+)', file_name)
    if match:
        extracted_part = match.group(1)
        print(extracted_part)  # 输出：540000
    else:
        print("未找到数字部分")

    # todo 将过滤的操作在这里进行，提取出需要计算的生态图，可节省计算开销
    # ecosystem_arrays = extract_ecosystem_arrays(tif_path)

    """
    将ecosystem_arrays替换成一个完整的生态图

    同时下面的
                for name, eco_array in filtered_ecosystem_arrays.items():
                替换成
                for name, eco_array in ecosystem_arrays.items():

    """
    import rasterio
    with rasterio.open(tif_path) as src:
        eco_total_data = src.read(1)

    # 将值为 255 的元素替换为 NaN
    eco_total_data = eco_total_data.astype(np.float32)  # 或者 np.float64
    eco_total_data[eco_total_data == 255] = np.nan
    ecosystem_arrays = {'全部生态系统': eco_total_data}

    # 遍历所有可用的公式并执行
    for formula_name in tqdm(available_formulas, desc="公式计算进度"):
        try:
            formula = next(f for f in formulas if f.name == formula_name)

            excel_parameter_names = [param["name"] for param in formula.required_params if param["is_excel"]]
            if excel_parameter_names:
                print("得到每个公式中是excel变量的变量名", excel_parameter_names)

            ecosystem_result_arrays = {}
            # 遍历每个生态图数组
            for name, eco_array in ecosystem_arrays.items():
                print(name)
                output_path = rf"{output_base_path}\{extracted_part}_{formula_name}_{name}_replaced_output.tif"
                # 检查文件是否存在
                if os.path.exists(output_path):
                    print(f"文件 {output_path} 已存在，跳过保存操作。")
                else:
                    # 通过替换eco.tif数据来输出tif
                    print(f"文件 {output_path} 不存在，保存操作。")
                    # 用excel转换换成的字典数据替换矩阵中的值
                    if excel_parameter_names:
                        # 如果有excel参数，将所有的excel参数取出并赋值
                        mapping_dict = {}
                        for param_name in excel_parameter_names:
                            # print("param_name", param_name)
                            # 得到excel参数的名字，用生态图乘以这个excel变量
                            print("Parameter参数", calculator.param_manager.get_param(param_name))
                            # print("Parameter参数值", calculator.param_manager.get_param(param_name).value)
                            mapping_dict = calculator.param_manager.get_param(param_name).value

                        formula_result = calculator.execute_formula(formula_name)
                        # todo 可以考虑直接在calculate_with_excel_mapping中传入excel参数对象
                        print("mapping_dict", mapping_dict)
                        ecosystem_result_arrays[name] = calculate_with_excel_mapping(eco_array, mapping_dict,
                                                                                     formula_result)

                    else:
                        print("没有需要excel的参数")
                        formula_result = calculator.execute_formula(formula_name)
                        ecosystem_result_arrays[name] = calculate_with_boolean_mapping(eco_array,
                                                                                       formula_result)  # formula_result = 0.099387

                    replace_tif_data_with_array(tif_path, ecosystem_result_arrays[name], output_path)

        except Exception as e:
            print(f"执行公式 {formula_name} 时出错: {e}")


if __name__ == '__main__':

    # 定义 JSON 文件夹路径
    json_folder_path = r'F:\code\dev\calc-gep-regulate-cqc\calc_v1\data\para_json_h'

    # 获取文件夹中的所有 JSON 文件
    json_files = [f for f in os.listdir(json_folder_path) if f.endswith('.json')]

    # 定义输出路径
    output_base_path_wrR = r'H:\GEP_result\2023\QwrR'
    output_base_path_wr = r'H:\GEP_result\2023\Qwr'
    os.makedirs(output_base_path_wr,exist_ok=True)


    # 使用 ProcessPoolExecutor 创建进程池
    with ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
        print("正在创建进程池...",os.cpu_count())
        # 提交任务到进程池
        futures = []
        for json_file in json_files:
            json_file_path = os.path.join(json_folder_path, json_file)
            print(f"正在提交文件: {json_file_path}")

            # 提交 wrR_variables 模式的任务
            # future_wrR = executor.submit(run_matrix_calculator, json_file_path, "wrR_variables", formulas=formulas_list_Q,
            #                             output_base_path=output_base_path_wrR)
            # futures.append(future_wrR)

            # 提交 wr_variables 模式的任务
            future_wr = executor.submit(run_matrix_calculator, json_file_path, "wr_variables", formulas=formulas_list_Q, output_base_path=output_base_path_wr)
            futures.append(future_wr)

        # 等待所有任务完成
        for future in as_completed(futures):
            try:
                result = future.result()
            except Exception as e:
                print(f"任务执行出错: {e}")
